#!/usr/bin/env python3

"""
"""


# built-ins
import pickle
from operator import itemgetter

# site-packages
import numpy as np
import matplotlib.pylab as plt
import matplotlib.cm as cm

from scipy import misc
from scipy import ndimage
from sklearn.cluster import k_means

# local-packages


def load_dataset(filename):
    """
    :return X: numpy.ndarray of shape (N, h, w) or (N, h, w, d)
    :return y: numpy.ndarray of shape N
    """
    with open(filename, 'rb') as pkl:
        data = pickle.load(pkl)
    X = np.array([row[0] for row in data])
    y = np.array([row[1] for row in data])
    return X, y


def get_output_size(input_size, recep_field_size, stride=1, padding=0):
    size = (input_size - recep_field_size + 2 * padding) / stride + 1
    if size == int(size):
        return int(size)
    else:
        raise KeyError('(W - F + 2P)/S + 1 must be integer, get %f' % size)


def im2col(im, blocksize, stride=1, padding=0):
    """
    Implementation of the MATLAB's im2col function.

    :type im: numpy.ndarray of shape (h, w)
    :param im: data matrix of image, may be padded

    :type blocksize: tuple of length 2 of int
    :param blocksize: receptive field size of (h, w)

    :rtype numpy.ndarray of shape (receptive-field ** 2 * d, col_ext * row_ext)
    :return columns from the image
    """
    h, w = im.shape
    fh, fw = blocksize
    col_ext = get_output_size(w, fw, stride, padding)
    row_ext = get_output_size(h, fh, stride, padding)
    start = np.arange(fh)[:, None] * w + np.arange(fw)
    offset = np.arange(row_ext)[:, None] * w + np.arange(col_ext)
    return np.take(im, start.ravel()[:, None] + offset.ravel())


def weight2row(weight):
    depth, h, w = weight.shape
    return weight.reshape((depth, h * w))


def cnn(im, _W, b):
    row = weight2row(_W)
    col = im2col(im, (5, 5))
    res = row.dot(col)
    out_size = get_output_size(im.shape[0], 5)
    assert res.shape[1] == out_size ** 2
    return res.reshape((res.shape[0], out_size, out_size))
